package SimpleLACore

import chisel3._
import chisel3.util._
import chisel3.util.random.LFSR

class SimpleCacheSRAM extends Module with SimpleLACoreParam {
  val io = IO(new Bundle {
    val core = new SimpleCacheIO
    val ram = Flipped(new SRAMIO)
  })
  require((nSets & (nSets - 1)) == 0, "nSets must be power of 2")
  require((lineWords & (lineWords - 1)) == 0, "wordsPerLine must be power of 2")
  require((nWays & (nWays - 1)) == 0, "nWays must be power of 2")
  
  def getIdx(addr: UInt):UInt = addr(idxBits + lgLineBytes - 1, lgLineBytes)
  def getTag(addr: UInt):UInt = addr(31, 32 - tagBits)
  def getOffset(addr: UInt):UInt = addr(lgLineBytes - 1, 2)

  val rand = LFSR(8)(lgnWays - 1, 0)

  val index = WireInit(UInt(idxBits.W), DontCare)
//  val metas = VecInit.tabulate(nWays)(_ => Mem(nSets, new SimpleCacheMeta)(index))
//  val datum = VecInit.tabulate(nWays)(_ => Mem(nSets, Vec(lineWords, Vec(4, UInt(8.W))))(index))
  val metas = Seq.fill(nWays)(Mem(nSets, new SimpleCacheMeta)(index))
  val datum = Seq.fill(nWays)(Mem(nSets, Vec(lineWords, Vec(4, UInt(8.W))))(index))
  val reset_counter = RegInit(0.U((idxBits + 1).W))
  val reset_done = reset_counter(idxBits)

  val coreBusyReg = RegInit(false.B)
  val coreReqReg = Reg(new SimpleCacheReq)
  val coreBusy = coreBusyReg || io.core.req.valid
  val coreReq = Mux(io.core.req.valid, io.core.req.bits, coreReqReg)

  when(io.core.resp.valid) {
    coreBusyReg := false.B
  }.elsewhen(io.core.req.valid) {
    coreBusyReg := true.B
    coreReqReg := io.core.req.bits
  }

  val buf = Reg(Vec(lineWords, Vec(4, UInt(8.W))))
  val bufBase = Reg(UInt(lgLineWords.W))
  val bufReqCnt = Reg(UInt(lgLineWords.W))
  val bufRespCnt = Reg(UInt(lgLineWords.W))
  val bufAddr = Reg(UInt((32 - lgLineBytes).W))

  val sReady :: sFill :: sSwap :: sWb :: Nil = Enum(4)
  val state = RegInit(sReady)
  val doResp = RegInit(false.B)
  val doReq = RegInit(false.B)

  io.ram.req := false.B
  io.ram.wr := DontCare
  io.ram.size := DontCare
  io.ram.addr := DontCare
  io.ram.wstrb := DontCare
  io.ram.wdata := DontCare
  io.core.resp.valid := false.B
  io.core.resp.bits := DontCare


  when(!reset_done){
    reset_counter := reset_counter + 1.U
    index := reset_counter
    metas.foreach(_.valid := false.B)
  }.otherwise{
    switch(state) {
      is(sReady) {
        when(coreBusy) {
          when(coreReq.cacop) {
            io.core.resp.valid := true.B
            switch(coreReq.code) {
              is(0.U) {
                index := getIdx(coreReq.addr)
                (metas zip UIntToOH(coreReq.addr(lgnWays - 1, 0)).asBools).foreach{ case (meta, sel) =>
                  when(sel){meta.valid := false.B}
                }
              }
              is(1.U) {
                index := getIdx(coreReq.addr)
                val waySel = UIntToOH(coreReq.addr(lgnWays - 1, 0)).asBools
                (metas zip waySel).foreach{ case (meta, hit) =>
                  when(hit){meta.valid := false.B}
                }
                val meta = Mux1H(waySel, metas)
                val data = Mux1H(waySel, datum)

                when(meta.dirty && meta.valid) {
                  state := sWb
                  doReq := true.B
                  buf := data//.asTypeOf(buf)
                  bufAddr := Cat(meta.tag, index)
                  bufBase := getOffset(coreReq.addr)
                  bufReqCnt := getOffset(coreReq.addr)
                  bufRespCnt := getOffset(coreReq.addr)
                }
              }
              is(2.U) {
                index := getIdx(coreReq.addr)
                val (dirty, data) = (datum zip metas).foldLeft((false.B, 0.U((lineWords * 32).W))) { case ((dirtyIn, dataIn), (line, meta)) =>
                  val hit = meta.valid && meta.tag === getTag(coreReq.addr)
                  (meta.dirty & hit | dirtyIn, line.asUInt & Fill(lineWords * 32, hit) | dataIn)
                }
                when(dirty){
                  state := sWb
                  doReq := true.B
                  buf := data.asTypeOf(buf)
                  bufAddr := coreReq.addr(31, lgLineBytes)
                  bufBase := getOffset(coreReq.addr)
                  bufReqCnt := getOffset(coreReq.addr)
                  bufRespCnt := getOffset(coreReq.addr)
                }
              }
            }
          }.elsewhen(coreReq.mat =/= 0.U){//CC
            val (hit, data) = (datum zip metas).foldLeft((false.B, 0.U(32.W))) { case ((hitIn, dataIn), (line, meta)) =>
              val hit = meta.valid && meta.tag === getTag(coreReq.addr)
              val data = line(getOffset(coreReq.addr))
              when(hit){
                (line(getOffset(coreReq.addr)) zip toBytes(coreReq.wdata) zip coreReq.wen.asBools).foreach{
                  case ((target, wData), wen) => when(wen){target := wData}
                }
                meta.dirty := coreReq.wen.orR
              }
              (hit | hitIn, data.asUInt & Fill(32, hit) | dataIn)
            }
            when(hit){
              io.core.resp.valid := !coreReq.preld
              io.core.resp.bits := data
            }.otherwise{
              state := sFill
              doReq := true.B
              doResp := !coreReq.preld
              bufAddr := coreReq.addr(31, lgLineBytes)
              bufBase := getOffset(coreReq.addr)
              bufReqCnt := getOffset(coreReq.addr)
              bufRespCnt := getOffset(coreReq.addr)
            }
          }.otherwise{//SUC
            when(!doResp){
              io.ram.req := true.B
              io.ram.wr := coreReq.wen.orR()
              io.ram.size := 2.U(2.W)
              io.ram.addr := coreReq.addr
              io.ram.wstrb := coreReq.wen
              io.ram.wdata := coreReq.wdata
            }
            when(io.ram.data_ok){
              doResp := false.B
              io.core.resp.valid := true.B
              io.core.resp.bits := io.ram.rdata
            }.elsewhen(io.ram.addr_ok){
              doResp := true.B
            }
          }
        }
      }
      is(sFill) {
        when(bufReqCnt =/= bufBase || doReq) {
          io.ram.req := true.B
          io.ram.wr := false.B
          io.ram.size := 2.U(2.W)
          io.ram.addr := Cat(bufAddr, bufReqCnt, 0.U(2.W))
          when(io.ram.addr_ok) {
            bufReqCnt := bufReqCnt + 1.U
            doReq := false.B
          }
        }
        when(io.ram.data_ok) {
          when(doResp){
            io.core.resp.valid := true.B
            io.core.resp.bits := io.ram.rdata
//            buf(bufRespCnt) := !coreReq.wMask & io.ram.rdata | coreReq.wMask & coreReq.wdata
            (io.ram.rBytes zip toBytes(coreReq.wdata) zip buf(bufRespCnt) zip coreReq.wen.asBools()).foreach{
              case (((rByte, wByte), bufByte), wen) =>
              bufByte := Mux(wen, wByte, rByte)
            }
            doResp := false.B
          }.otherwise{
            buf(bufRespCnt) := io.ram.rBytes
          }
          bufRespCnt := bufRespCnt + 1.U
          when(bufRespCnt + 1.U === bufBase) {
            state := sSwap
          }
        }
      }
      is(sSwap){
        val replaceWay = rand // random replacement
        index := bufAddr // scala will truncate tag bits
        val waySel = UIntToOH(replaceWay).asBools
        (datum zip waySel).foreach{ case (line, sel) =>
          when(sel){
            line := buf//.asTypeOf(line)
          }
        }
        (metas zip waySel).foreach{ case (meta, sel) =>
          when(sel){
            meta.tag := bufAddr(tagBits + idxBits - 1, idxBits)
            meta.dirty := false.B
            meta.valid := true.B
          }
        }
        buf := Mux1H(waySel, datum)//.asTypeOf(buf)
        val meta = Mux1H(waySel, metas)
        bufAddr := Cat(meta.tag, index)
        when(meta.dirty && meta.valid) {
          state := sWb
          doReq := true.B
          //we already have bufCnt === bufBase here
        }.otherwise {
          state := sReady
        }
      }
      is(sWb) {
        when(bufReqCnt =/= bufBase || doReq) {
          io.ram.req := true.B
          io.ram.wr := true.B
          io.ram.size := 2.U(2.W)
          io.ram.addr := Cat(bufAddr, bufReqCnt)
          io.ram.wstrb := 0xf.U
          io.ram.wdata := buf(bufReqCnt).asUInt
          when(io.ram.addr_ok) {
            bufReqCnt := bufReqCnt + 1.U
            doReq := false.B
          }
        }
        when(io.ram.data_ok) {
          bufRespCnt := bufRespCnt + 1.U
          when(bufRespCnt + 1.U === bufBase) {
            state := sReady
          }
        }
      }
    }
  }
}